import numpy as np
import matplotlib.pyplot as plt
import itk
import os


def get_file_list(folder, extension):
    file_list = []
    file_name_list = []
    for dir_path, dir_names, file_names in os.walk(folder):
        for file in file_names:
            file_type = os.path.splitext(file)[1]
            if file_type == extension:
                file_fullname = os.path.join(dir_path, file)
                file_list.append(file_fullname)
    return file_list


if __name__ == "__main__":

    folder = '/mnt/520/haoshl/MC-GPU/output/medical_2/proj_chest/scale_1.0/'
    out_folder = '/mnt/520/haoshl/RTK_recon/data/proj_chest/scale_1.0/'
    file_list = get_file_list(folder, '.raw')
    file_list = sorted(file_list)

    total_proj = np.zeros(shape=(len(file_list), 237, 324))
    primary_proj = np.zeros(shape=(len(file_list), 237, 324))
    scatter_proj = np.zeros(shape=(len(file_list), 237, 324))

    for i in range(len(file_list)):
        raw_data = np.fromfile(file_list[i], dtype='float32')
        raw_data = np.reshape(raw_data, (2, 237, 324))
        total = raw_data[0, :, :]
        scatter = raw_data[1, :, :]
        total = np.flipud(total)
        scatter = np.flipud(scatter)
        primary = total - scatter

        total_proj[i, :, :] = total
        primary_proj[i, :, :] = primary
        scatter_proj[i, :, :] = scatter

    I0_total = np.max(total_proj)
    total_proj = -np.log(np.divide(total_proj, I0_total))
    total_proj[total_proj == np.inf] = 0
    total_proj[total_proj == np.nan] = 0

    I0_primary = np.max(primary_proj)
    primary_proj = -np.log(np.divide(primary_proj, I0_primary))
    primary_proj[primary_proj == np.inf] = 0
    primary_proj[primary_proj == np.nan] = 0

    for i in range(len(file_list)):

        total = total_proj[i, :, :]
        total = itk.image_from_array(total)
        total = total.astype(itk.F)
        itk.imwrite(total, out_folder + 'total/' + '{:0>4d}'.format(i) + '.tif')

        primary = primary_proj[i, :, :]
        primary = itk.image_from_array(primary)
        primary = primary.astype(itk.F)
        itk.imwrite(primary, out_folder + 'primary/' + '{:0>4d}'.format(i) + '.tif')

        # scatter = scatter_proj[i, :, :]
        # scatter = itk.image_from_array(scatter)
        # scatter = scatter.astype(itk.F)
        # itk.imwrite(scatter, out_folder + 'scatter/' + '{:0>4d}'.format(i) + '.tif')